import torch
import torch.nn as nn

from test import PointEdgeReg, StandardScaler

# ==== 还原你训练好的模型和 scaler ====
ckpt = torch.load("checkpoints/best_reg.ckpt", map_location="cpu", weights_only=False)

core = PointEdgeReg(in_dim=4, hidden=128, out_dim=1, k=16)   # 和训练时一致
core.load_state_dict(ckpt["model_state"])
scaler = StandardScaler.from_dict(ckpt["scaler"])

# ==== 包一层，把标准化放进 forward 里，C# 只喂原始特征 ====
class WrappedModel(nn.Module):
    def __init__(self, core, mean, std):
        super().__init__()
        self.core = core
        self.register_buffer("mean", mean)
        self.register_buffer("std",  std)
    def forward(self, x):
        # x: [N,4]  (x,y,z,strength)
        x = (x - self.mean) / self.std
        return self.core(x)

wrapped = WrappedModel(core,
                       torch.from_numpy(scaler.mean.numpy()).float().squeeze(0),
                       torch.from_numpy(scaler.std.numpy()).float().squeeze(0))
wrapped.eval()

# ==== 导出 TorchScript ====
ts_model = torch.jit.script(wrapped)   # 用 script，别 trace
ts_model.save("model_ts.pt")
print("saved TorchScript to model_ts.pt")
